import torch.nn as nn
import torch
import numpy as np
from collections import deque
import random

class ReplayBuffer:
    def __init__(self, buffer_size):
        self.buffer = deque(maxlen=buffer_size)
        self.buffer_size = buffer_size

    def add(self, state, action, reward, next_state, done):
        state, state_pos = state
        state = [torch.squeeze(state, dim=0), torch.squeeze(state_pos, dim=0)]
        next_state, next_state_pos = next_state
        next_state = [torch.squeeze(next_state, dim=0), torch.squeeze(next_state_pos, dim=0)]

        data = tuple([*state, action, reward, *next_state, done])
        self.buffer.append(data)

    def sample(self, batch_size):
        batch = random.sample(self.buffer, batch_size)
        state, state_pos, action, reward, next_state, next_state_pos, done = map(torch.stack, zip(*batch))
        state = [state, state_pos]
        next_state = [next_state, next_state_pos]
        return state, action, reward, next_state, done

    def is_full(self):
        return len(self.buffer) == self.buffer_size

    def __len__(self):
        return len(self.buffer)
